from template import *
from collect import collect


import sys

itime = -1
if len(sys.argv) == 2:
	t = int(sys.argv[1])
elif len(sys.argv) > 2:
	print("error: should have one argument, as time step")
else:
	t = itime



ref_path = "conduction_limit_re0.2/data"
ref_config = json.load( open(ref_path+"/parameter.inp", "r") )
species_e = ref_config["Species"][0]
species_D1 = ref_config["Species"][1]
print(species_e["mass"], species_D1["mass"])

vdf_position = ref_config["Diagnostic"][0]["vedf_position"]
print("vdf_position: ", vdf_position)
vdf_index = []
for position in vdf_position:
    index = int(position / ref_config["cell_length"][0])
    vdf_index.append(index)
print("vdf_index: ", vdf_index)
cell_point_number = int(ref_config["sim_length"][0] / ref_config["cell_length"][0])
print("cell_point_number: ", cell_point_number)


vdf_velocity_array_e = collect("/Diagnostic/", "vdf_velocity_array", path = ref_path)[t, 0, 0, :]
vdf_velocity_array_D1 = collect("/Diagnostic/", "vdf_velocity_array", path = ref_path)[t, 0, 1, :]

print("size of x: ", vdf_velocity_array_e.shape)

#============ define V_T thermal velocity to normalize velocity 
Te = 50.0
V_Te = math.sqrt( 2.0 * const.e * Te / species_e["mass"] )

TD1 = 50.0
V_TD1 = math.sqrt( 2.0 * const.e * Te / species_D1["mass"] )

xmin = vdf_velocity_array_e.min()
xmax = vdf_velocity_array_e.max()

list_fe_sheath = []
list_fe_mid = []
list_fe_source = []
list_fD1_sheath = []
list_fD1_mid = []
list_fD1_source = []

list_label = [r"$C \mathrm{_r=0.2}$", r"$C \mathrm{_r=0.4}$", r"$C \mathrm{_r=0.6}$", r"$C \mathrm{_r=0.8}$", r"$C \mathrm{_r=0.95}$"]

list_xlabel = [r"$v_{\parallel} (v_{T_{\mathrm{e}}})$", r"$v_{\parallel} (v_{T_{\mathrm{D^+}}})$", r"$v_{\parallel} (v_{T_{\mathrm{e}}})$", r"$v_{\parallel} (v_{T_{\mathrm{D^+}}})$", r"$v_{\parallel} (v_{T_{\mathrm{e}}})$", r"$v_{\parallel} (v_{T_{\mathrm{D^+}}})$"]

list_ylabel = [r"$f \mathrm{_e(source)}$", r"$f \mathrm{_{D^+}(source)}$",r"$f \mathrm{_e(mid)}$", r"$f \mathrm{_{D^+}(mid)}$", r"$f \mathrm{_e(sheath)}$", r"$f \mathrm{_{D^+}(sheath)}$"]

list_path = ["conduction_limit_re0.2/data", "conduction_limit_re0.4/data", "conduction_limit_re0.6/data", "conduction_limit_re0.8/data", "conduction_limit_re0.95/data"]
for path_item in list_path:
    val = collect("/Diagnostic/", "vdf_parallel_to_B", path = path_item)[t, :, :, :]

    list_fe_sheath.append(val[0, 0, :])
    list_fe_mid.append(val[3, 0, :])
    list_fe_source.append(val[5, 0, :])

    list_fD1_sheath.append(val[0, 1, :])
    list_fD1_mid.append(val[3, 1, :])
    list_fD1_source.append(val[5, 1, :])



##inite the fig of matplotlib
fig=plt.figure(figsize=(15,10))
fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.35)

##============fe source======================================================
ax=fig.add_subplot(3,2,1)

for i in range(len(list_fe_source)):
    line0=ax.plot(vdf_velocity_array_e/V_Te, list_fe_source[i], label = list_label[i])

ax.set_xlabel(list_xlabel[0], fontsize = label_fontsize)
#ax.set_ylim(bottom = 0.0)
ax.set_ylabel(list_ylabel[0], fontsize = label_fontsize)
#ax.grid(True)
ax.set_yscale('log') 

num1 = 0
num2 = 1.4
num3 = 2
num4 = 0
handles, labels = ax.get_legend_handles_labels()
ax.legend(flip(handles, 3), flip(labels, 3), framealpha = 0.2, bbox_to_anchor=(num1, num2), loc=num3, borderaxespad=num4, ncol = 3)
ax.text(-0.1, -0.15, r"$\mathbf{(a)}$", transform=ax.transAxes)

##============fD1 source======================================================
ax=fig.add_subplot(3,2,2)

for i in range(len(list_fD1_source)):
    line0=ax.plot(vdf_velocity_array_D1/V_TD1, list_fD1_source[i], label = list_label[i])

ax.set_xlabel(list_xlabel[1], fontsize = label_fontsize)
#ax.set_ylim(bottom = 0.0)
ax.set_ylabel(list_ylabel[1], fontsize = label_fontsize)
#ax.grid(True)
ax.set_yscale('log') 
#ax.legend(framealpha = 0.2)
ax.text(-0.1, -0.15, r"$\mathbf{(b)}$", transform=ax.transAxes)

##============ fe mid ======================================================
ax=fig.add_subplot(3,2,3)

for i in range(len(list_fe_mid)):
    line0=ax.plot(vdf_velocity_array_e/V_Te, list_fe_mid[i], label = list_label[i])

ax.set_xlabel(list_xlabel[2], fontsize = label_fontsize)
#ax.set_ylim(bottom = 0.0)
ax.set_ylabel(list_ylabel[2], fontsize = label_fontsize)
#ax.grid(True)
ax.set_yscale('log') 
#ax.legend(framealpha = 0.2)
ax.text(-0.1, -0.15, r"$\mathbf{(c)}$", transform=ax.transAxes)



##============ fD1 mid ======================================================
ax=fig.add_subplot(3,2,4)

for i in range(len(list_fD1_mid)):
    line0=ax.plot(vdf_velocity_array_D1/V_TD1, list_fD1_mid[i], label = list_label[i])

ax.set_xlabel(list_xlabel[3], fontsize = label_fontsize)
#ax.set_ylim(bottom = 0.0)
ax.set_ylabel(list_ylabel[3], fontsize = label_fontsize)
#ax.grid(True)
ax.set_yscale('log') 
#ax.legend(framealpha = 0.2)
ax.text(-0.1, -0.15, r"$\mathbf{(d)}$", transform=ax.transAxes)


##============ fe sheath ======================================================
ax=fig.add_subplot(3,2,5)

for i in range(len(list_fe_sheath)):
    line0=ax.plot(vdf_velocity_array_e/V_Te, list_fe_sheath[i], label = list_label[i])

ax.set_xlabel(list_xlabel[4], fontsize = label_fontsize)
#ax.set_ylim(bottom = 0.0)
ax.set_ylabel(list_ylabel[4], fontsize = label_fontsize)
#ax.grid(True)
ax.set_yscale('log') 
#ax.legend(framealpha = 0.2)
ax.text(-0.1, -0.15, r"$\mathbf{(e)}$", transform=ax.transAxes)


##============ fD1 sheath ======================================================
ax=fig.add_subplot(3,2,6)

for i in range(len(list_fD1_sheath)):
    line0=ax.plot(vdf_velocity_array_D1/V_TD1, list_fD1_sheath[i], label = list_label[i])

ax.set_xlabel(list_xlabel[5], fontsize = label_fontsize)
#ax.set_ylim(bottom = 0.0)
ax.set_ylabel(list_ylabel[5], fontsize = label_fontsize)
#ax.grid(True)
ax.set_yscale('log') 
#ax.legend(framealpha = 0.2)
ax.text(-0.1, -0.15, r"$\mathbf{(f)}$", transform=ax.transAxes)


fig_file_name = "Vdf_all_cl_" + str(t) + ".svg"
fig.savefig(fig_file_name, dpi = 300)
##fig.show()       #when the program finishes,the figure disappears
#plt.axis('equal')
#plt.show()         #The command is OK
